package main

import (
	"context"
	"fmt"
	"log"

	"cloud.google.com/go/bigquery"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

// BigQueryServer 处理BigQuery请求的服务器实现
type BigQueryServer struct {
	connections map[string]*BigQueryHandler
}

// BigQueryConnection BigQuery连接配置
type BigQueryConnection struct {
	Project         string `json:"project"`
	Dataset         string `json:"dataset"`
	CredentialsFile string `json:"credentials_file,omitempty"`
}

// NewBigQueryServer 创建新的BigQuery服务器
func NewBigQueryServer(connections map[string]BigQueryConnection) (*BigQueryServer, error) {
	server := &BigQueryServer{
		connections: make(map[string]*BigQueryHandler),
	}

	// 初始化所有BigQuery连接
	for name, cfg := range connections {
		ctx := context.Background()
		handler, err := NewBigQueryHandler(ctx, cfg.Project, cfg.Dataset, cfg.CredentialsFile)
		if err != nil {
			return nil, fmt.Errorf("无法连接到BigQuery '%s': %v", name, err)
		}
		server.connections[name] = handler
		log.Printf("BigQuery连接 '%s' 已初始化", name)
	}

	return server, nil
}

// HandleRequest 处理BigQuery请求
func (s *BigQueryServer) HandleRequest(req SQLRequest) SQLResponse {
	// 获取BigQuery处理器
	handler, exists := s.connections[req.Connection]
	if !exists {
		return SQLResponse{Error: "BigQuery连接不存在: " + req.Connection}
	}

	// BigQuery不支持事务
	if req.Operation == "begin" || req.Operation == "commit" || req.Operation == "rollback" {
		return SQLResponse{Error: "BigQuery不支持事务"}
	}

	// 根据操作类型执行不同的处理
	switch req.Operation {
	case "query":
		return s.handleQuery(handler, req)
	case "exec":
		return s.handleExec(handler, req)
	case "prepare":
		return s.handlePrepare(handler, req)
	default:
		return SQLResponse{Error: "未知操作类型: " + req.Operation}
	}
}

// 处理BigQuery查询
func (s *BigQueryServer) handleQuery(handler *BigQueryHandler, req SQLRequest) SQLResponse {
	// BigQuery不支持事务，忽略RequestID
	if req.RequestID != "" {
		return SQLResponse{Error: "BigQuery不支持事务"}
	}

	// 执行查询
	it, err := handler.Query(req.SQL, req.Args)
	if err != nil {
		return SQLResponse{Error: err.Error()}
	}

	// 读取所有行
	rows, columns, err := handler.ReadAllRows(it)
	if err != nil {
		return SQLResponse{Error: err.Error()}
	}

	return SQLResponse{
		Columns: columns,
		Rows:    rows,
	}
}

// 处理BigQuery执行
func (s *BigQueryServer) handleExec(handler *BigQueryHandler, req SQLRequest) SQLResponse {
	// BigQuery不支持事务，忽略RequestID
	if req.RequestID != "" {
		return SQLResponse{Error: "BigQuery不支持事务"}
	}

	// 执行命令
	rowsAffected, err := handler.Exec(req.SQL, req.Args)
	if err != nil {
		return SQLResponse{Error: err.Error()}
	}

	return SQLResponse{
		Rows:    [][]interface{}{{0, rowsAffected}}, // LastInsertId始终为0
		Columns: []string{"last_insert_id", "rows_affected"},
	}
}

// 处理BigQuery预处理语句（模拟）
func (s *BigQueryServer) handlePrepare(handler *BigQueryHandler, req SQLRequest) SQLResponse {
	// BigQuery不支持预处理语句，但我们可以存储SQL以便后续使用
	stmtID := fmt.Sprintf("bq:%s", req.SQL)

	return SQLResponse{
		Rows:    [][]interface{}{{stmtID}},
		Columns: []string{"stmt_id"},
	}
}

// BigQueryHandler BigQuery数据库处理器
type BigQueryHandler struct {
	client  *bigquery.Client
	project string
	dataset string
	ctx     context.Context
}

// NewBigQueryHandler 创建BigQuery处理器
func NewBigQueryHandler(ctx context.Context, project, dataset, credentialsFile string) (*BigQueryHandler, error) {
	var client *bigquery.Client
	var err error

	if credentialsFile != "" {
		client, err = bigquery.NewClient(ctx, project, option.WithCredentialsFile(credentialsFile))
	} else {
		client, err = bigquery.NewClient(ctx, project)
	}

	if err != nil {
		return nil, err
	}

	return &BigQueryHandler{
		client:  client,
		project: project,
		dataset: dataset,
		ctx:     ctx,
	}, nil
}

// Close 关闭连接
func (h *BigQueryHandler) Close() error {
	if h.client != nil {
		return h.client.Close()
	}
	return nil
}

// Query 执行查询并返回结果
func (h *BigQueryHandler) Query(sql string, args []interface{}) (*bigquery.RowIterator, error) {
	// 创建查询
	q := h.client.Query(sql)

	// 设置查询参数（如果有）
	if len(args) > 0 {
		params := make([]bigquery.QueryParameter, len(args))
		for i, arg := range args {
			params[i] = bigquery.QueryParameter{
				Name:  h.getParamName(i),
				Value: arg,
			}
		}
		q.Parameters = params
	}

	// 执行查询
	job, err := q.Run(h.ctx)
	if err != nil {
		return nil, err
	}

	// 等待查询完成
	status, err := job.Wait(h.ctx)
	if err != nil {
		return nil, err
	}

	if err := status.Err(); err != nil {
		return nil, err
	}

	// 获取结果
	return job.Read(h.ctx)
}

// Exec 执行DML命令并返回影响的行数
func (h *BigQueryHandler) Exec(sql string, args []interface{}) (int64, error) {
	// 创建查询
	q := h.client.Query(sql)

	// 设置查询参数（如果有）
	if len(args) > 0 {
		params := make([]bigquery.QueryParameter, len(args))
		for i, arg := range args {
			params[i] = bigquery.QueryParameter{
				Name:  h.getParamName(i),
				Value: arg,
			}
		}
		q.Parameters = params
	}

	// 执行查询
	job, err := q.Run(h.ctx)
	if err != nil {
		return 0, err
	}

	// 等待查询完成
	status, err := job.Wait(h.ctx)
	if err != nil {
		return 0, err
	}

	if err := status.Err(); err != nil {
		return 0, err
	}

	// 获取DML影响的行数
	var rowsAffected int64 = 0
	if dmlStats, ok := status.Statistics.Details.(*bigquery.QueryStatistics); ok && dmlStats != nil {
		rowsAffected = dmlStats.NumDMLAffectedRows
	}

	return rowsAffected, nil
}

// ReadAllRows 从迭代器读取所有行
func (h *BigQueryHandler) ReadAllRows(it *bigquery.RowIterator) ([][]interface{}, []string, error) {
	// 提取列名
	schema := it.Schema
	columns := make([]string, len(schema))
	for i, field := range schema {
		columns[i] = field.Name
	}

	// 读取所有行
	var rows [][]interface{}
	for {
		var values []bigquery.Value
		err := it.Next(&values)
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, nil, err
		}

		// 转换为接口数组
		row := make([]interface{}, len(values))
		for i, v := range values {
			row[i] = v
		}
		rows = append(rows, row)
	}

	return rows, columns, nil
}

// getParamName 获取参数名称
func (h *BigQueryHandler) getParamName(index int) string {
	return fmt.Sprintf("p%d", index)
}
